import os

# 【关键修改】在代码中强制设置 Hugging Face 镜像
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

import logging
from typing import List, Dict, Any
from dotenv import load_dotenv

# --- 标准做法：将所有import放在文件顶部 ---
from langchain_community.document_loaders import TextLoader, PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings # 【关键修改】导入HuggingFace嵌入模型

# 从配置文件和工具文件导入
from utils.config import TEXT_SPLITTER_CONFIG, RETRIEVER_CONFIG, SUPPORTED_FORMATS
from utils.file_utils import validate_document




# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class RAGSystem:
    """
    一个基于LangChain的检索增强生成（RAG）系统。
    它封装了。
    """

    def __init__(self):
        """初始化RAG系统，设置API并加载模型。"""
        load_dotenv()

        # 从环境变量中获取API配置
        self.api_key = os.getenv("DEEPSEEK_API_KEY")
        self.api_base = "https://api.deepseek.com/v1"

        if not self.api_key:
            raise ValueError("错误：请在 .env 文件中设置 DEEPSEEK_API_KEY")

        self.llm = ChatOpenAI(
            model="deepseek-chat",
            openai_api_key=self.api_key,
            openai_api_base=self.api_base,
            temperature=0.1
        )
        # 【关键修改】嵌入模型使用本地的HuggingFace模型
        print("🔄 正在加载本地嵌入模型，首次运行需要下载模型，请稍候...")
        self.embeddings = HuggingFaceEmbeddings(
            model_name="BAAI/bge-large-zh-v1.5",  # 使用顶级的中文嵌入模型
            model_kwargs={'device': 'cpu'},  # 如果你有GPU，可以改为 'cuda'
            encode_kwargs={'normalize_embeddings': True}  # 建议开启，对于bge模型
        )
        print("✅ 嵌入模型加载完成！")

        self.vector_store = None
        self.qa_chain = None

        logger.info("RAG系统初始化成功，模型已加载。")

    def load_documents(self, file_path: str) -> List[Any]:
        """
        根据文件扩展名加载单个文档（支持.txt和.pdf）。
        在加载前会验证文件。
        """
        logger.info(f"正在加载文档: {file_path}")

        # 【关键修改】使用工具函数进行文件验证
        if not validate_document(file_path):
            raise FileNotFoundError(f"文档验证失败: {file_path}")

        file_extension = os.path.splitext(file_path)[1].lower()
        if file_extension not in SUPPORTED_FORMATS:
            raise ValueError(f"不支持的文件格式: {file_extension}。支持的格式: {SUPPORTED_FORMATS}")

        if file_extension == '.txt':
            loader = TextLoader(file_path, encoding='utf-8')
        elif file_extension == '.pdf':
            loader = PyPDFLoader(file_path)
        else:
            # 理论上不会走到这里，因为前面已经验证
            loader = None

        documents = loader.load()
        logger.info(f"成功加载 {len(documents)} 个文档片段。")
        return documents

    def split_documents(self, documents: List[Any]) -> List[Any]:
        """
        将文档分割成小块，配置从config.py读取。
        """
        logger.info("正在分割文档...")
        # 【关键修改】使用配置文件中的参数
        text_splitter = RecursiveCharacterTextSplitter(**TEXT_SPLITTER_CONFIG)

        chunks = text_splitter.split_documents(documents)
        logger.info(f"文档已分割为 {len(chunks)} 个文本块。")
        return chunks

    def create_vector_store(self, chunks: List[Any], save_path: str = None):
        """根据文本块创建向量存储，并可选择保存到本地。"""
        logger.info("正在创建向量存储...")
        self.vector_store = FAISS.from_documents(chunks, self.embeddings)

        if save_path:
            self.vector_store.save_local(save_path)
            logger.info(f"向量存储已保存至: {save_path}")

        # 使用配置文件中的参数
        self.retriever = self.vector_store.as_retriever(**RETRIEVER_CONFIG)
        logger.info("向量存储与检索器创建成功。")

    def load_vector_store(self, load_path: str):
        """从本地加载已有的向量存储。"""
        logger.info(f"正在从 {load_path} 加载向量存储...")
        self.vector_store = FAISS.load_local(
            load_path,
            self.embeddings,
            allow_dangerous_deserialization=True
        )
        self.retriever = self.vector_store.as_retriever(**RETRIEVER_CONFIG)
        logger.info("向量存储与检索器加载成功。")

    def create_qa_chain(self):
        """构建用于问答的LCEL链。"""
        if not self.retriever:
            raise ValueError("检索器未初始化，请先创建或加载向量存储。")

        logger.info("正在创建问答链...")

        prompt_template = """
        你是一个专业的问答助手。请根据以下提供的上下文来回答用户的问题。
        你的回答应该简洁、准确，并且只基于上下文信息。
        如果上下文中没有相关信息，请回答“根据提供的上下文，我无法回答该问题。”

        上下文:
        {context}

        问题:
        {question}
        """
        prompt = ChatPromptTemplate.from_template(prompt_template)

        self.qa_chain = (
                {"context": self.retriever, "question": RunnablePassthrough()}
                | prompt
                | self.llm
                | StrOutputParser()
        )
        logger.info("问答链创建成功。")

    def ask_question(self, question: str) -> str:
        """使用问答链来回答问题。"""
        if not self.qa_chain:
            raise ValueError("问答链未创建，请先调用 create_qa_chain() 方法。")

        logger.info(f"收到问题: {question}")
        response = self.qa_chain.invoke(question)
        logger.info("已生成答案。")
        return response

    def search_similar_documents(self, query: str, k: int = 3) -> List[Dict]:
        """检索与查询最相似的文档片段。"""
        if not self.retriever:
            raise ValueError("检索器未初始化，请先创建或加载向量存储。")

        docs = self.vector_store.similarity_search(query, k=k)
        results = [
            {"rank": i + 1, "content": doc.page_content, "metadata": doc.metadata}
            for i, doc in enumerate(docs)
        ]
        return results